import numpy as np

def pauli_z_i(i, nqubits):
    """
        return pauli z operate on i-th qubit
    """
    sigma_z = np.array([[1, 0], [0, -1]])
    I = np.array([[1, 0], [0, 1]])
    
    op_string = []
    for j in range(nqubits):
        if j != i:
            op_string.append(I)
        else:
            op_string.append(sigma_z)
    
    res = op_string[0]
    for op in op_string[1:]:
        res = np.kron(res, op)
    
    return res

def get_ham_matrix(BIG_vertex_weight_matrix, BIG_adjacent_matrix, P=6):
    n_l = len(BIG_vertex_weight_matrix)
    n_p = len(BIG_vertex_weight_matrix[0])
    
    nqubits = n_l * n_p

    ham = np.zeros((2**nqubits, 2**nqubits))
    w_list = BIG_vertex_weight_matrix.flatten()
    
    for i in range(nqubits):
        for j in range(nqubits):
            if BIG_adjacent_matrix[i, j] == 0 and i!=j:
                ham += P/2*(pauli_z_i(i, nqubits)@pauli_z_i(j, nqubits)-pauli_z_i(i, nqubits)-pauli_z_i(j, nqubits))
    
    for i in range(nqubits):
        ham += w_list[i]*pauli_z_i(i, nqubits)
    
    return ham


def get_ham_matrix_from_circuit(circ, theta):
    return circ.matrix(pr=theta)
